# bert_head.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class BertClassificationHead(nn.Module):
    """
    用于从 BERT Encoder 输出得到分类结果。
    """
    def __init__(self, hidden_size=768, num_classes=2, dropout=0.1, use_cls=True):
        super().__init__()
        self.use_cls = use_cls
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(hidden_size, num_classes)

    def forward(self, encoder_output):
        """
        参数:
            encoder_output: [seq_len, hidden_dim] → 来自 Voltage 推理系统的输出

        返回:
            logits: [num_classes]
        """
        if self.use_cls:
            cls_rep = encoder_output[0]  # [CLS] token at index 0
        else:
            cls_rep = encoder_output.mean(dim=0)  # Mean Pooling

        rep = self.dropout(cls_rep)
        logits = self.linear(rep)
        return logits
